//go:build go1.20
// +build go1.20

/*
Copyright 2024 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package test

import (
	"context"
	"errors"
	"regexp"
	"strings"
	"testing"

	"cloud.google.com/go/spanner"
	"cloud.google.com/go/spanner/internal"
	stestutil "cloud.google.com/go/spanner/internal/testutil"
	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/sdk/trace/tracetest"
	"google.golang.org/api/option"
	"google.golang.org/grpc"
)

func TestSpannerTracesWithOpenTelemetry(t *testing.T) {
	ctx := context.Background()
	te := newOpenTelemetryTestExporter(false, true)

	t.Cleanup(func() {
		te.Unregister(ctx)
	})

	minOpened := uint64(1)
	server, client, teardown := setupMockedTestServerWithConfig(t, spanner.ClientConfig{
		SessionPoolConfig: spanner.SessionPoolConfig{
			MinOpened: minOpened,
		},
	}, option.WithGRPCDialOption(grpc.WithStreamInterceptor(otelgrpc.StreamClientInterceptor(otelgrpc.WithTracerProvider(te.tp)))))
	defer teardown()

	waitFor(t, func() error {
		if isMultiplexEnabled {
			if uint64(server.TestSpanner.TotalSessionsCreated()) == minOpened+1 {
				return nil
			}
		}
		if uint64(server.TestSpanner.TotalSessionsCreated()) == minOpened {
			return nil
		}
		return errors.New("not yet initialized")
	})

	const (
		transactionTag = "tx-tag"
		requestTag     = "req-tag"
	)
	_, err := client.ReadWriteTransactionWithOptions(context.Background(), func(ctx context.Context, tx *spanner.ReadWriteTransaction) error {
		iter := tx.QueryWithOptions(ctx, spanner.NewStatement(stestutil.SelectSingerIDAlbumIDAlbumTitleFromAlbums), spanner.QueryOptions{
			RequestTag: requestTag,
		})
		return iter.Do(func(r *spanner.Row) error {
			return nil
		})
	}, spanner.TransactionOptions{
		TransactionTag: transactionTag,
	})
	if err != nil {
		t.Fatalf("ReadWriteTransaction failed: %v", err)
	}
	waitFor(t, func() error {
		if findSpan(te.spans(), "cloud.google.com/go/spanner.ReadWriteTransactionWithOptions") {
			return nil
		}
		return errors.New("readWriteTransactionWithOptions span not found")
	})
	waitFor(t, func() error {
		if findSpan(te.spans(), "cloud.google.com/go/spanner.RowIterator") {
			return nil
		}
		return errors.New("rowIterator span not found")
	})
	// Wait until all spans have been collected.
	// The span that is generated by the gRPC interceptor is only ended
	// when the trailers have been received, which happens asynchronously.
	waitFor(t, func() error {
		if findSpan(te.spans(), "google.spanner.v1.Spanner/ExecuteStreamingSql") {
			return nil
		}
		return errors.New("gRPC span not found")
	})

	spans := te.spans()
	if len(spans) == 0 {
		t.Fatal("No spans were exported")
	}
	for _, spanName := range []string{"cloud.google.com/go/spanner.Query", "cloud.google.com/go/spanner.RowIterator", "google.spanner.v1.Spanner/ExecuteStreamingSql"} {
		span := getSpan(spans, spanName)
		if span == nil {
			t.Fatalf("Failed to find span %q", spanName)
		}
		if spanName == "google.spanner.v1.Spanner/ExecuteStreamingSql" {
			if g, w := span.Status.Code, codes.Ok; g != w {
				t.Errorf("span %q status code mismatch\n Got: %v\nWant: %v", spanName, g, w)
			}
			continue
		}
		if g, w := span.Status.Code, codes.Unset; g != w {
			t.Errorf("span %q status code mismatch\n Got: %v\nWant: %v", spanName, g, w)
		}
	}

	txAttributes := map[attribute.Key]string{
		"db.name":             "[DATABASE]",
		"instance.name":       "[INSTANCE]",
		"cloud.region":        "global",
		"gcp.client.version":  internal.Version,
		"gcp.client.repo":     "googleapis/google-cloud-go",
		"gcp.client.artifact": "cloud.google.com/go/spanner",
		"transaction.tag":     transactionTag,
	}
	var foundTaggedQuerySpan bool

	for _, span := range spans {
		if !strings.HasPrefix(span.Name, "cloud.google.com/go/spanner.") {
			continue
		}
		if g, w := span.Status.Code, codes.Unset; g != w {
			t.Errorf("span %q status code mismatch\n Got: %v\nWant: %v", span.Name, g, w)
		}

		switch span.Name {
		case "cloud.google.com/go/spanner.ReadWriteTransactionWithOptions":
			for key, value := range txAttributes {
				if !hasAttribute(&span, key, value) {
					t.Errorf("span %q was missing attribute %q with value %q", span.Name, key, value)
				}
			}
		case "cloud.google.com/go/spanner.RowIterator":
			if hasAttribute(&span, "statement.tag", requestTag) {
				foundTaggedQuerySpan = true
				if !hasAttribute(&span, "db.statement", stestutil.SelectSingerIDAlbumIDAlbumTitleFromAlbums) {
					t.Errorf("tagged span %q was missing db.statement", span.Name)
				}
				if !hasAttribute(&span, "transaction.tag", transactionTag) {
					t.Errorf("tagged span %q was missing transaction.tag", span.Name)
				}
			}
		}
	}
	if !foundTaggedQuerySpan {
		t.Errorf("did not find a RowIterator span with the specified statement.tag %q", requestTag)
	}
}

func TestSpanAnnotatedWithXGoogSpannerRequestID_unary(t *testing.T) {
	targetSpanName := "cloud.google.com/go/spanner.BatchCreateSessions"
	if isMultiplexEnabled {
		targetSpanName = "cloud.google.com/go/spanner.CreateSession"
	}
	testSpanAnnotatedWithXGoogSpannerRequestID(t, targetSpanName)
}

func TestSpanAnnotatedWithXGoogSpannerRequestID_streaming(t *testing.T) {
	testSpanAnnotatedWithXGoogSpannerRequestID(t, "cloud.google.com/go/spanner.RowIterator")
}

func testSpanAnnotatedWithXGoogSpannerRequestID(t *testing.T, targetSpanName string) {
	ctx := context.Background()
	te := newOpenTelemetryTestExporter(false, true)

	t.Cleanup(func() {
		te.Unregister(ctx)
	})

	minOpened := uint64(1)
	server, client, teardown := setupMockedTestServerWithConfig(t, spanner.ClientConfig{
		SessionPoolConfig: spanner.SessionPoolConfig{
			MinOpened: minOpened,
		},
	})
	defer teardown()

	waitFor(t, func() error {
		if isMultiplexEnabled {
			if uint64(server.TestSpanner.TotalSessionsCreated()) == minOpened+1 {
				return nil
			}
		}
		if uint64(server.TestSpanner.TotalSessionsCreated()) == minOpened {
			return nil
		}
		return errors.New("not yet initialized")
	})

	iter := client.Single().Query(context.Background(), spanner.NewStatement(stestutil.SelectSingerIDAlbumIDAlbumTitleFromAlbums))
	iter.Do(func(r *spanner.Row) error {
		return nil
	})
	spans := te.spans()
	if len(spans) == 0 {
		t.Fatal("No spans were exported")
	}

	// Find the target span
	var targetSpan tracetest.SpanStub
	found := false
	for _, span := range spans {
		if span.Name == targetSpanName {
			targetSpan = span
			found = true
			break
		}
	}

	if !found {
		t.Fatalf("Could not find target span with name: %q", targetSpanName)
	}

	var headerKey attribute.Key = "x_goog_spanner_request_id"
	ourHeaderFound := false
	foundValue := ""
	for _, attrKV := range targetSpan.Attributes {
		if attrKV.Key == headerKey {
			foundValue = attrKV.Value.AsString()
			ourHeaderFound = foundValue != ""
			break
		}
	}

	if !ourHeaderFound {
		t.Fatalf("did not find our header: %q", headerKey)
	}

	// It must match the desired format of x-goog-spanner-request-id.
	var reg = regexp.MustCompile("\\d+\\.([a-z0-9]{16})(\\.\\d+){4}")
	if !reg.MatchString(foundValue) {
		t.Fatalf("Regex=%q did not match %q", reg, foundValue)
	}
}

func findSpan(spans tracetest.SpanStubs, spanName string) bool {
	return getSpan(spans, spanName) != nil
}

func getSpan(spans tracetest.SpanStubs, spanName string) *tracetest.SpanStub {
	for _, span := range spans {
		if span.Name == spanName {
			return &span
		}
	}
	return nil
}

func hasAttribute(span *tracetest.SpanStub, key attribute.Key, value string) bool {
	for _, attr := range span.Attributes {
		// Special case for cloud.region, value depends on the environment.
		if attr.Key == "cloud.region" {
			return true
		}
		if attr.Key == key && attr.Value.AsString() == value {
			return true
		}
	}
	return false
}
